import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
from tensorflow.keras.datasets import fashion_mnist

# 加载Fashion MNIST数据集
(x_train_fashion, y_train_fashion), (x_test_fashion, y_test_fashion) = fashion_mnist.load_data()

# 数据预处理 - 与MNIST相同
x_test_fashion = x_test_fashion.reshape(-1, 28, 28, 1).astype('float32') / 255.0

# 定义Fashion MNIST的类别名称
#class_names = ['T恤/上衣', '裤子', '套衫', '连衣裙', '外套', '凉鞋', '衬衫', '运动鞋', '包', '短靴']
class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']
# 加载TFLite模型
tflite_model_file = "./mnist_tflite_models/mnist_model.tflite"
interpreter = tf.lite.Interpreter(model_path=str(tflite_model_file))
interpreter.allocate_tensors()

# 获取输入和输出张量信息
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 使用TFLite模型进行推理的函数
def predict_fashion(image):
    # 调整图像形状以匹配模型输入
    input_data = np.expand_dims(image, axis=0)
    
    # 设置输入张量
    interpreter.set_tensor(input_details[0]['index'], input_data)
    
    # 运行推理
    interpreter.invoke()
    
    # 获取输出张量
    output_data = interpreter.get_tensor(output_details[0]['index'])
    return np.argmax(output_data), np.max(output_data)


# 随机选择几个样本进行测试
num_samples = 5
indices = np.random.choice(len(x_test_fashion), num_samples)

plt.figure(figsize=(15, 6))
for i, idx in enumerate(indices):
    plt.subplot(1, num_samples, i+1)
    
    # 获取测试图像
    test_image = x_test_fashion[idx]
    
    # 预测
    predicted_class, confidence = predict_fashion(test_image)
    
    # 显示图像和预测结果
    plt.imshow(test_image.reshape(28, 28), cmap='gray')
    plt.title(f'Prediction: {class_names[predicted_class]}\nActual: {class_names[y_test_fashion[idx]]}\nConfidence: {confidence:.2f}')
    plt.axis('off')

plt.tight_layout()
plt.show()

# 评估模型在Fashion MNIST上的整体性能
correct_predictions = 0
total_predictions = len(x_test_fashion)

for i in range(total_predictions):
    predicted_class, _ = predict_fashion(x_test_fashion[i])
    if predicted_class == y_test_fashion[i]:
        correct_predictions += 1

accuracy = correct_predictions / total_predictions
print(f"在Fashion MNIST上的准确率: {accuracy:.4f}")
